import warnings
from joblib import Parallel, delayed
from GradientBoostingRegressor import GradientBoostingModel
from SVM import SupportVectorMachineModel
from Ridge import RidgeRegressionModel

warnings.filterwarnings("ignore", category=UserWarning)



def train_model(model_class, task_type, model_type, excel_path, features, target, **kwargs):
    model_instance = model_class(task_type, model_type, excel_path, features, target, **kwargs)
    results = model_instance.train()
    model_instance.plot_loss_curve()  # 绘制损失曲线
    return results

def main():
    # 输入参数
    excel_path = r'C:\Users\13945\Desktop\MLDS铝合金成分设计数据.xlsx'
    features = ['Si/％', 'Mn/％', 'Zn/％', 'Mg/％', 'Cu/％', 'Cr/％', 'Zr/％', 'Ti/％', 'Fe/％', 'Ni/％', 'other']
    target = 'Ultimate tensile strength/MPa'

    # 定义模型参数
    params_gb = {
        'task_type': 'regression',
        'model_type': 'GradientBoosting',
        'test_size': 0.2,
        'random_state': 42,
        'learning_rate': 0.1,
        'max_depth': 6,
        'min_samples_leaf': 1,
        'min_samples_split': 2,
        'n_estimators': 90
    }

    params_svr = {
        'task_type': 'regression',
        'model_type': 'SVR',
        'test_size': 0.2,
        'random_state': 42,
        'kernel': 'rbf',
        'C': 1.0,
        'epsilon': 0.1,
        'gamma': 'scale',
        'max_iter': -1
    }

    params_rid = {
        'task_type': 'regression',
        'model_type': 'Ridge',
        'test_size': 0.2,
        'random_state': 42,
        'alpha': 1.0
    }

    tasks = [
        delayed(train_model)(GradientBoostingModel, **params_gb, excel_path=excel_path, features=features,
                             target=target),
        delayed(train_model)(SupportVectorMachineModel, **params_svr, excel_path=excel_path, features=features,
                             target=target),
        delayed(train_model)(RidgeRegressionModel, **params_rid, excel_path=excel_path, features=features,
                             target=target)
    ]

    results = Parallel(n_jobs=-1)(tasks)

    # 打印结果
    for result in results:
        print(result)

if __name__ == "__main__":
    main()